package com.maiko.maikoaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;


/**
 * 为AI对话系统提供基于文件持久化的记忆功能。用于持久化存储每个对话的历史消息
 * 使用 Kryo 序列化框架将消息列表保存为 .kryo 文件，每个对话一个文件。
 * @author: Maiko7
 * @create: 2025-06-15 9:15
 */

// 这个类实现了ChatMemory接口，用于存储AI与用户的历史对话消息，每个对话会被保存到一个本地文件中。
public class FileBasedChatMemory implements ChatMemory {

    // BASE_DIR：所有对话文件保存的目录路径
    private final String BASE_DIR;
    /**
     * Kryo 是一个高性能 Java 对象序列化工具，用于将对象写入文件或从文件中读出。
     * 此处静态初始化 Kryo 实例，用于整个类内部复用。
     */
    private static final Kryo kryo = new Kryo();

    static {
        // 允许 Kryo 序列化未注册的类，避免必须注册类型
        kryo.setRegistrationRequired(false);
        // 设置实例化策略，确保无构造器类也能被反序列化
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    /**
     * 构造方法，初始化存储目录路径。
     *
     * @param dir 本地对话文件存储的根目录
     */
    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        /**
         * Q：这不就是创建一个文件吗？
         * A：这只是创建了一个表示“路径”的 File 对象，既可以是目录，也可以是文件，但它本身不创建任何实际的文件或目录。
         * baseDir.mkdirs(); 这才是真正的创建目录
         */
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();  // 如果目录不存在，则创建
        }
    }


    /**
     * 将一批消息追加到指定会话中，并保存到本地文件。如果该会话已有历史消息，将在原有基础上追加。
     * 调用流程：
     * 1. 从文件中加载已有的消息；
     * 2. 把新消息加入列表；
     * 3. 序列化整个消息列表到文件中。
     * @param conversationId 一个对话唯一标识（比如用户ID、sessionID）。
     * @param messages 要添加的新消息列表（List<Message>），例如：用户输入和 AI 回复。
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        // 第一步：读取或创建历史消息
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        // 第二步：追加新消息
        conversationMessages.addAll(messages);
        // 第三步：保存文件
        saveConversation(conversationId, conversationMessages);
    }



    /**
     * 根据会话 ID，从本地文件中获取最近 lastN 条对话记录。
     * @param conversationId  对话的唯一标识
     * @param lastN 获取的消息数量
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 获取全部历史消息
        List<Message> allMessages = getOrCreateConversation(conversationId);
        // 提取后 N 条消息
        return allMessages.stream()
                /**
                 * .skip(n)跳过前面n条消息，只保留最后N条
                 * allMessages.size()	当前会话总共有多少条消息
                 * allMessages.size() - lastN	需要跳过的消息数量，跳过前面的
                 * Math.max(0, ...) 如果消息数量不足 lastN，就跳过 0 条，防止负数
                 */
                .skip(Math.max(0, allMessages.size() - lastN)) // 跳过前面部分，只取后面 N 条
                .toList();  // 把Stream转换为List返回
    }


    /**
     * 用于清除指定对话的历史记录。
     * @param conversationId    要清除的对话 ID，对应一段存储在磁盘上的对话历史。
     */
    @Override
    public void clear(String conversationId) {
        /**
         * 作用是根据对话 ID 获取保存对话内容的文件对象。
         * 实际生成的文件名形如：BASE_DIR/chat123.kryo
         * 假设 BASE_DIR = "data/"
         * getConversationFile("chat123") → File("data/chat123.kryo")
         */
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            // 如果对应的.kryo 文件存在，就执行delete()删除操作。
            file.delete();
        }
    }


    /**
     * 根据会话 ID 加载本地的历史消息列表。如果文件不存在或读取失败，则返回空列表。
     *
     * 思路：它这个代码怎么写的呢？你把InMemory的代码发给GPT，然后说用Kryo实现即可。
     *
     * @param conversationId 会话 ID
     * @return 历史消息列表（可能为空）
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        // 调用 getConversationFile(conversationId) 方法，根据会话 ID 返回对应的本地文件路径。
        File file = getConversationFile(conversationId);
        // 如果本地没有历史消息文件，或读取失败，就返回这个空列表。
        List<Message> messages = new ArrayList<>();
        /**
         * 判断 .kryo 文件是否存在；
         * 使用 Kryo 反序列化文件内容，转换为 Java 对象（ArrayList<Message>）；
         * kryo.readObject(...) 会读取整个对象结构，并返回反序列化后的 messages；
         * 如果读取过程中出错（比如文件损坏），则打印异常，继续返回空列表。
         */
        if (file.exists()) {
            /**
             * Input input = new Input(new FileInputStream(file))用来读取一个本地文件的二进制内容，用于后面的反序列化
             * 创建一个 Kryo 的输入流，用于从指定的文件中读取数据，之后可配合 kryo.readObject(...) 把字节内容还原为 Java 对象。
             * 注意Input是导入这个的import com.esotericsoftware.kryo.io.Input;
             */
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class); // 反序列化
            } catch (IOException e) {
                e.printStackTrace(); // 文件损坏或反序列化异常
            }
        }
        return messages;
    }

    /**
     * 用于将指定会话的消息列表序列化并保存到本地文件
     * @param conversationId    对话的唯一标识
     * @param messages  要保存的消息列表
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        // 获取对应的文件路径
        File file = getConversationFile(conversationId);
        // 打开文件输出流并序列化（序列化：将内存中的Java对象转换为字节流，写入文件）保存
        try (Output output = new Output(new FileOutputStream(file))) {
            // 将messages对象序列化写入文件
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取当前会话对应的本地文件路径。
     * 文件名格式为：{BASE_DIR}/{conversationId}.kryo
     *
     * @param conversationId 会话 ID
     * @return File 对象，表示会话记录文件
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }

}
